import numpy as np
import torch

def load_txt(file_path,batch_size=1):
    data = np.loadtxt(file_path, delimiter=' ', dtype=np.float32)
    tensor_np = np.reshape(data, (batch_size, -1, 3))
    tensor = torch.from_numpy(tensor_np)
    tensor.cuda(0)
    return tensor

def   save_txt_(name,tensor):
    file_path = f'result_py/{name}.txt'
    tensor = tensor.squeeze()
    if tensor.ndim == 2 and tensor.shape[0] < tensor.shape[1]:
        tensor = tensor.T
    tensor=tensor.cpu().numpy()
    np.savetxt(file_path, tensor, fmt='%.6f', delimiter=' ', newline='\n')

def   save_txt_fc(name,tensor):
    file_path = f'result_py/{name}.txt'
    tensor=tensor.cpu().numpy()
    with open(file_path, 'w') as file:
        for batch in tensor:
            for batch_str in batch:
                file.write(format(batch_str, ".6f") + "\n")
                
def  save_txt_fc3(name,tensor):
    file_path = f'result_py/{name}.txt'
    tensor=tensor.cpu().numpy()
    with open(file_path, 'w') as file:
        for batch in tensor:
            for row in batch:
                for row_str in row:
                    file.write(format(row_str, ".6f") + "\n")
    
def  save_txt(name, tensor):
    tensor = tensor.cpu().numpy()
    shape = tensor.shape
    num_dims = len(shape)
    if num_dims == 1:
        file_path = f'result_py/{name}.txt'
        np.savetxt(file_path, tensor, fmt='%.6f', delimiter=' ', newline='\n')
    elif num_dims == 2:
        if shape[0] < shape[1]:
            tensor = tensor.T
        file_path = f'result_py/{name}.txt'
        np.savetxt(file_path, tensor, fmt='%.6f', delimiter=' ', newline='\n')
    elif tensor.ndim == 3:
        batch_size = tensor.shape[0]
        file_path = f'result_py/{name}.txt'  # 注意这里只有一个文件名
        if shape[1] < shape[2]:
            tensor = tensor.transpose((0, 2, 1))
        with open(file_path, 'a') as f:  # 以追加模式打开文件
            for i in range(batch_size):
                np.savetxt(f, tensor[i], fmt='%.6f', delimiter=' ', newline='\n')
    else:
        raise ValueError("Unsupported number of dimensions. Tensor must have 1, 2, or 3 dimensions.")